# -*- coding: utf-8 -*-
"""
@Env 
@Time 2024/9/9 下午2:52
@Author yzpang
@Function: 模型训练服务
"""
import os

from ..entity.request import TrainMessage
from modelserver.configs.model_config import MODEL_ROOT_PATH, DATASET_ROOT_PATH
from modelserver.configs.base_config import get_logger
from modelserver.train.train import train_model
from ..config.registry import train_model_registry

from multiprocessing import Process


def train_start(train_info: TrainMessage):
    """启动训练"""
    model_name = train_info.modelName
    if model_name in train_model_registry:
        raise KeyError(f"模型{model_name}训练中...")
    model_path = train_info.modelPath
    label_path = os.path.join(DATASET_ROOT_PATH, train_info.labelPath)
    data_path = os.path.join(DATASET_ROOT_PATH, train_info.dataPath)
    eval_path = os.path.join(DATASET_ROOT_PATH, train_info.evalPath)
    output_path = os.path.join(MODEL_ROOT_PATH, train_info.outputPath)
    classification_type = train_info.classificationType

    kwargs = {}
    kwargs["label_path"] = label_path
    kwargs["train_path"] = data_path
    kwargs["eval_path"] = eval_path
    kwargs["output_dir"] = output_path
    kwargs["classification_type"] = classification_type
    kwargs.update(train_info.extra)
    get_logger().info(kwargs)

    process = Process(target=train_model,
                      name='model_train_' + model_name,
                      args=(model_path,),
                      kwargs=kwargs)
    process.start()
    train_model_registry[model_name] = process
    return f"模型{model_name}启动训练"


def train_stop(train_info: TrainMessage):
    """停止训练"""
    model_name = train_info.modelName
    if model_name not in train_model_registry:
        raise KeyError(f"模型{model_name}未训练")
    process = train_model_registry[model_name]
    process.kill()
    del train_model_registry[model_name]
    return f"模型{model_name}停止训练"


def train_list():
    """训练模型列表"""
    for key in list(train_model_registry.keys()):
        if not check_train_model(key):
            del train_model_registry[key]

    return list(train_model_registry.keys())


def check_train_model(model_name: str):
    """ 检查模型进程是否有效 """
    process = train_model_registry[model_name]
    return process.is_alive()
